# -*- coding: utf-8 -*-
# @File:     generate_task.py
# @Author:
# @DateTime: 2025/06/06/10:48
import datetime
import json
import time
from typing import Optional, Union, Any, List, Dict
from fastapi.encoders import jsonable_encoder
from sqlalchemy import or_
from sqlalchemy.orm import Session

from app.api.generate.utils.generate_utils import crud_generate_record
from app.api.generate.utils.redis_utils import key_add_order_queue
from app.api.generate.utils.rabbitmq_queue import queue_send_message
from app.conf import settings
from app.db.database import get_db
from app.models.system import Users
from app.utils import security
from app.utils.custom_exc import CustomException
from app.utils.custom_response import success_response, error_response
from fastapi import APIRouter, Query, Body, Path, Depends, UploadFile, Form, File, Request
from pydantic import BaseModel, Field

from app.models.generate import GenerateTask, MenuList, UserMembership, GenerateResult, GenerateParamImage
from app.utils.curd_base import ModelCRUD
from app.utils.file_storage import get_storage_file, save_storage_file, delete_storage_file
from app.utils.utils import get_now_time_string

router = APIRouter()


class CRUDGenerateTask(ModelCRUD):

    model = GenerateTask


class GenerateTaskReqData(BaseModel):
    menu_id: Optional[int] = Field(default=None, description='菜单id')
    input_data: Union[dict, list] = Field(description='输入参数')
    priority: Optional[int] = Field(default=None, description='优先级')


class CreateGenerateTask(BaseModel):
    user_id: int = Field(description='用户')
    menu_id: Optional[int] = Field(default=None, description='菜单id')
    params: Optional[str] = Field(default=None, description='参数')

    send_queue: Optional[str] = Field(default=None, description='发送队列')
    order_queue: Optional[str] = Field(default=None, description='排序队列')


class ResGenerateTask(CreateGenerateTask):
    error_describe: Optional[str] = Field(default=None, description='错误描述')


class DeleteModel(BaseModel):
    id: list[int] = Field(description='ID列表')


async def auth_user_generate_number(cnn, membership, generate_point, acquire_timeout=45, lock_key='generate_acquire', db=None):
    end = time.time() + acquire_timeout
    while time.time() < end:
        if not await cnn.get(lock_key):
            await cnn.set(lock_key, lock_key, acquire_timeout)
            try:
                user_able_num = membership.total - membership.use_number
                if user_able_num >= generate_point:
                    # membership.use_number = membership.use_number + generate_point
                    # db.add(membership)
                    # db.commit()
                    # db.refresh(membership)
                    return True
                else:
                    return False
            except:
                CustomException(err_desc='生成失败')
            finally:
                await cnn.delete(lock_key)
        time.sleep(0.5)
    raise CustomException(err_desc='生成失败')


@router.post("/generate_task", summary="生成任务", response_model=ResGenerateTask, name='生成任务')
async def user_membership(
    request: Request,
    *,
    db: Session = Depends(get_db),
    req_data: GenerateTaskReqData,
    current_user: Users = Depends(security.verify_access_token),
) -> Any:

    user_id: int = int(current_user.id)
    crud = CRUDGenerateTask()

    input_params = req_data.input_data

    menu = db.query(MenuList).filter(MenuList.id == req_data.menu_id, MenuList.enable == 1).first()
    if not menu:
        return error_response(message='未找到该应用')

    now_time = datetime.datetime.now()
    membership = db.query(UserMembership).filter(UserMembership.user_id == user_id, UserMembership.start_time <= now_time, UserMembership.expire_time>=now_time).first()
    if not membership:
        membership = db.query(UserMembership).filter(UserMembership.start_time == None, UserMembership.expire_time == None, UserMembership.user_id == int(user_id)).first()  # type: ignore
        if not membership:
            return error_response(message='你没有开通套餐')


    # 验证并行
    queryset = db.query(GenerateTask).filter(GenerateTask.user_id == user_id, GenerateTask.state == 1)
    exce_num = queryset.count()
    # print(exce_num)
    if exce_num >= membership.task_num:
        return error_response(message='已达到最大并发数')

    cnn = request.app.state.redis
    if not await auth_user_generate_number(cnn, membership, menu.point, db=db):
        return error_response(message='生成失败生成点数不足')

    task_data = {
        # 'task_uuid': get_now_time_string(),
        'user_id': current_user.id,
        'menu_id': req_data.menu_id,
        'params': json.dumps(req_data.input_data, indent=4, ensure_ascii=False),
        'send_queue': menu.queue,
        'order_queue': menu.order,
    }

    instance = crud.create(db, data=CreateGenerateTask.validate(task_data))
    only = ('id', 'task_uuid', 'menu_id', 'state', 'params', 'error_describe', 'create_time')
    data = instance.to_dict(only=only, datetime_format='%Y-%m-%d %H:%M:%S')

    record = crud_generate_record(instance, db)

    await cnn.set(f'task_{instance.id}', json.dumps(data), 300)

    send_data = input_params
    send_data['task_id'] = instance.id
    send_data['workflowId'] = menu.workflowId
    send_data['resp'] = {'user_id': user_id, 'membership': membership.id, 'send_queue': menu.queue,
                         'order_queue': menu.order, 'record_id': record.id, 'point': menu.point}

    priority = req_data.priority
    send_priority = membership.priority
    if priority and priority <= membership.priority:
        send_priority = priority
    queue_send_message(send_data, menu.queue, priority=send_priority)

    await key_add_order_queue(menu.order, instance.id)

    return success_response(data=data)


able_filter_list = [
    {'name': 'user_id', 'source_filed': 'user_id', 'lookup_expr': 'eq', 'label': '等于user', 'type': int, 'method': None, 'required': True},
    {'name': 'menu_id', 'source_filed': 'menu_id', 'lookup_expr': 'eq', 'label': '等于menu_id', 'type': int, 'method': None, 'required': False},
]


@router.get("/generate_task", summary="获取生成任务", response_model=ResGenerateTask, name='获取生成任务')
async def user_membership(
    *,
    db: Session = Depends(get_db),
    page: int = Query(1, description='页码', ge=1),
    page_size: int = Query(10, description='页面大小', le=50, ge=1),
    menu_id: int = Query(None, description='菜单ID'),
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    order_by_fields = ('-id',)
    crud = CRUDGenerateTask()

    req_data = {
        'user_id': current_user.id,
    }
    if menu_id:
        req_data['menu_id'] = menu_id

    queryset, total = crud.get_page_queryset(db, page=page, page_size=page_size, req_data=req_data, able_filter_list=able_filter_list, order_by_fields=order_by_fields)
    only = ('id', 'task_uuid', 'menu_id', 'state', 'params', 'error_describe', 'create_time')
    result_only = ('id', 'path', 'file_type', 'download', 'create_time')
    li = []
    for instance in queryset:
        data = instance.to_dict(only=only, datetime_format='%Y-%m-%d %H:%M:%S')
        data['menu_name'] = instance.menu.name if instance.menu else None
        result_queryset = db.query(GenerateResult).filter(GenerateResult.task_id == int(instance.id))
        result_li = []
        for result_ins in result_queryset:
            result_data = result_ins.to_dict(only=result_only, datetime_format='%Y-%m-%d %H:%M:%S')
            result_data['path_url'] = get_storage_file(result_ins.path)
            result_li.append(result_data)

        data['result_list'] = result_li
        li.append(data)
    resp = {
        "list": li,
        'total': total
    }
    return success_response(data=resp)


able_filter_list1 = [
    {'name': 'user_id', 'source_filed': 'user_id', 'lookup_expr': 'eq', 'label': '等于user', 'type': int, 'method': None, 'required': True},
    {'name': 'id', 'source_filed': 'id', 'lookup_expr': 'eq', 'label': '等于id', 'type': int, 'method': None, 'required': True},
]


@router.get("/generate_task/{instance_id}", summary="获取生成任务", response_model=ResGenerateTask, name='获取生成任务')
async def user_membership(

    *,
    db: Session = Depends(get_db),
    instance_id,
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    crud = CRUDGenerateTask()

    req_data = {
        'user_id': current_user.id,
        'id': instance_id
    }
    queryset = crud.get_all_queryset(db, req_data=req_data, able_filter_list=able_filter_list1).first()
    if not queryset:
        return error_response(message="未找到数据")
    only = ('id', 'task_uuid', 'menu_id', 'state', 'params', 'error_describe', 'create_time')
    items = crud.serializer(queryset, only=only, datetime_format='%Y-%m-%d %H:%M:%S')
    resp = {
        "list": items
    }
    return success_response(data=resp)


class SelectGenerateStateReqData(BaseModel):
    id: list[int] = Field(description='ID数组')


@router.post("/get_task_state", summary="获取生成任务状态", response_model=ResGenerateTask, name='获取生成任务状态')
async def user_membership(
    request: Request,
    *,
    db: Session = Depends(get_db),
    req_data: SelectGenerateStateReqData,
    current_user: Users = Depends(security.verify_access_token),
) -> Any:
    crud = CRUDGenerateTask()

    cnn = request.app.state.redis

    id_list = req_data.id
    id_list = id_list[:5]
    li = []

    for instance_id in id_list:

        state_data = await cnn.get(f'task_{instance_id}')

        if state_data:
            li.append(json.loads(state_data))
        else:
            req_data = {
                'user_id': current_user.id,
                'id': instance_id
            }
            instance = crud.get_all_queryset(db, req_data=req_data, able_filter_list=able_filter_list1).first()
            if not instance:
                continue
            only = ('id', 'task_uuid', 'menu_id', 'state', 'params', 'error_describe', 'create_time')
            data = instance.to_dict(only=only, datetime_format='%Y-%m-%d %H:%M:%S')
            result_li = []
            if instance.state == 3:
                result_only = ('id', 'path', 'file_type', 'download', 'create_time')
                result_queryset = db.query(GenerateResult).filter(GenerateResult.task_id == int(instance.id))

                for result_ins in result_queryset:
                    result_data = result_ins.to_dict(only=result_only, datetime_format='%Y-%m-%d %H:%M:%S')
                    result_data['path_url'] = get_storage_file(result_ins.path)
                    result_li.append(result_data)
            data['result_list'] = result_li
            li.append(data)
            await cnn.set(f'task_{instance.id}', json.dumps(data), 300)

    return success_response(data=li)


# -------------------------------------------------------------------

able_filter_list_back = [
    {'name': 'user_id', 'source_filed': 'user_id', 'lookup_expr': 'eq', 'label': '等于user', 'type': int, 'method': None, 'required': False},
    {'name': 'menu_id', 'source_filed': 'menu_id', 'lookup_expr': 'eq', 'label': '等于menu_id', 'type': int, 'method': None, 'required': False},
]

@router.get("/manage_generate_task", summary="后台获取生成任务", response_model=ResGenerateTask, name='后台获取生成任务')
async def user_membership(
    *,
    db: Session = Depends(get_db),
    page: int = Query(1, description='页码', ge=1),
    page_size: int = Query(10, description='页面大小', le=50, ge=1),
    menu_id: int = Query(None, description='菜单ID'),
    user_id: int = Query(None, description='用户ID'),
    current_user: Users = Depends(security.verify_admin_permission),
) -> Any:
    order_by_fields = ('-id',)
    crud = CRUDGenerateTask()

    req_data = {}
    if user_id:
        req_data['user_id'] = user_id
    if menu_id:
        req_data['menu_id'] = menu_id

    queryset, total = crud.get_page_queryset(db, page=page, page_size=page_size, req_data=req_data, able_filter_list=able_filter_list_back, order_by_fields=order_by_fields)
    only = ('id', 'task_uuid', 'menu_id', 'state', 'params', 'error_describe', 'create_time')
    result_only = ('id', 'path', 'file_type', 'download', 'create_time')
    li = []
    for instance in queryset:
        data = instance.to_dict(only=only, datetime_format='%Y-%m-%d %H:%M:%S')
        data['menu_name'] = instance.menu.name if instance.menu else None
        result_queryset = db.query(GenerateResult).filter(GenerateResult.task_id == int(instance.id))
        result_li = []
        for result_ins in result_queryset:
            result_data = result_ins.to_dict(only=result_only, datetime_format='%Y-%m-%d %H:%M:%S')
            result_data['path_url'] = get_storage_file(result_ins.path)
            result_li.append(result_data)

        data['result_list'] = result_li
        li.append(data)
    resp = {
        "list": li,
        'total': total
    }
    return success_response(data=resp)


class SelectModel(BaseModel):
    id: int = Field(description='ID列表')
    key: str = Field(description='链接key')

able_filter_list111 = [
    {'name': 'id', 'source_filed': 'id', 'lookup_expr': 'eq', 'label': '等于id', 'type': int, 'method': None, 'required': True},
]

@router.post("/select_task", summary="验证生成任务", response_model=ResGenerateTask, name='验证生成任务')
async def user_membership(

    *,
    db: Session = Depends(get_db),
    req_data: SelectModel,
    # current_user: Users = Depends(security.verify_access_token),
) -> Any:
    crud = CRUDGenerateTask()

    if req_data.key != settings.RABBITMQ_CLIENT_KEY:
        return error_response(message='连接密钥不正确')

    s_data = {
        'id': req_data.id
    }
    queryset = crud.get_all_queryset(db, req_data=s_data, able_filter_list=able_filter_list111).first()
    if not queryset:
        return error_response(message="未找到数据")
    only = ('id', 'task_uuid', 'state', )
    items = queryset.to_dict(only=only)
    return success_response(data=items)

'''
{
  "menu_id": 6,
  "input_data": {
    "user_data": [
        {
            "name": "图片",
            "nodeId": "20",
            "fieldName": "image",
            "fieldValue": "https://linkdesign-resource.com.cn/images/exam_image/task_eg/300/cloth_img_1748573817611.png"
        },
        {
            "name": "文本描述",
            "nodeId": "21",
            "fieldName": "text",
            "fieldValue": "模特摆出拍照姿势，展示衣服的高端与时尚感"
        },
        {
            "name": "最大宽度",
            "nodeId": "34",
            "fieldName": "value",
            "fieldValue": "768"
        }
    ],
    "result_node": [
        "11"
    ],
    "workflowId": "1934890416552484865",
    "nodeInfoList": [
        {
            "nodeId": "13",
            "fieldName": "seed",
            "fieldValue": "12345"
        }
    ]
},
  "params": {}
}




{
    "user_data": [
        {
            "name": "图片",
            "type": "image",
            "nodeId": "1",
            "fieldName": "image",
            "fieldValue": ""
        },
        {
            "name": "蒙版",
            "type": "mask",
            "link": "146",
            "nodeId": "147",
            "fieldName": "image",
            "fieldValue": ""
        },
        {
            "name": "文本描述",
            "type": "text",
            "nodeId": "147",
            "fieldName": "prompt",
            "fieldValue": ""
        }
    ],
    "result_node": ["142"],
    "nodeInfoList": [
        {
            "nodeId": "145",
            "fieldName": "seed",
            "fieldValue": "12345"
        }
    ]
}
'''






